import torch
import matplotlib.pyplot as plt
import numpy as np

# 读取数据
data = torch.load("dataset.pt", weights_only=False)

print("数据类型:", type(data))
print("样本数量:", len(data))

# 取一个样本（可修改索引）
sample = data[250]
print("样本类型:", type(sample))
# print("样本内容:", sample)
src = sample["src"] if isinstance(sample, dict) else sample[0]
tgt = sample["tgt"] if isinstance(sample, dict) else sample[1]

src = np.array(src)
tgt = np.array(tgt)

print("src shape:", src.shape)  # (10, 9)
print("tgt shape:", tgt.shape)  # (10, 1)

# 绘制每个特征及其目标值
num_features = src.shape[1]
time_steps = np.arange(src.shape[0])

plt.figure(figsize=(12, 14))
for i in range(num_features):
    ax1 = plt.subplot(5, 2, i + 1)  # 5行2列布局

    # 左轴：src
    ax1.plot(time_steps, src[:, i], 'b-o', label=f"src feature {i}", linewidth=1.5, markersize=4)
    ax1.set_xlabel("Timestep")
    ax1.set_ylabel("Src Value", color='b')
    ax1.tick_params(axis='y', labelcolor='b')
    ax1.grid(True, linestyle='--', alpha=0.6)

    # 右轴：tgt
    ax2 = ax1.twinx()
    if tgt.shape[1] == 1:
        ax2.plot(time_steps, tgt[:, 0], 'r--s', label="tgt", linewidth=1.5, markersize=4)
    else:
        ax2.plot(time_steps, tgt[:, i], 'r--s', label=f"tgt feature {i}", linewidth=1.5, markersize=4)
    ax2.set_ylabel("Tgt Value", color='r')
    ax2.tick_params(axis='y', labelcolor='r')

    # 标题
    ax1.set_title(f"Feature {i}")

plt.tight_layout()
plt.show()
